import machine_lib as ml
import pandas as pd
s = ml.login()
region = "EUR"  # EUR TWN
id = "model135"  # "other423" analyst48

if region == "AMR":
    uni = "TOP600"
elif region == "JPN":
    uni = "TOP1600"
elif region == "USA":
    uni = "TOP3000"
elif region == "ASI":
    uni = "MINVOL1M"
elif region == "KOR":
    uni = "TOP600"
elif region == "TWN":
    uni = "TOP500"
elif region == "HKG":
    uni = "TOP800"
elif region == "GLB":
    uni = "TOP3000"
elif region == "CHN":
    uni = "TOP2000U"
elif region == "EUR":
    uni = "TOP1200"
datafields = ml.get_datafields(s, dataset_id = id, region=region, universe=uni)
print(datafields)

# datafields = datafields[(datafields["coverage"] > 0.6)]
#print(datafields)

# pc_fields = process_datafields(datafields, "matrix")

pc_fields = ml.process_datafields(datafields, "vector")

def custom_shuffle(alpha_list, interval):
    shuffled_list = []
    n = len(alpha_list)
    for i in range(interval):
        shuffled_list.append(alpha_list[i:n:interval])
    flattened_list = [item for sublist in shuffled_list for item in sublist]
    return flattened_list

import os
operators = ml.arsenal + ml.ts_ops
# print(operators)
first_order = ml.first_order_factory(pc_fields, operators)

name = f"{region}_fo_{id}"
print(name)
init_decay = 5
fo_alpha_list = []
for alpha in first_order:
    fo_alpha_list.append((alpha, init_decay))

print(f"len of fo_alpha_list: {len(fo_alpha_list)}")

# shuffled_list = custom_shuffle(fo_alpha_list, 129)

df_fo_alpha = pd.DataFrame(fo_alpha_list, columns=["alpha", "decay"])
output_dir = "./fo_alpha_list/"
os.makedirs(output_dir, exist_ok=True)
try:
    df_fo_alpha.to_csv(f"./fo_alpha_list/{name}.csv", index=False)
    print(f"DataFrame successfully saved.")
except Exception as e:
    print(f"Error saving DataFrame to CSV: {e}")